package org.dwpd.utils;

import dev.langchain4j.data.message.*;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.bson.types.ObjectId;
import org.dwpd.domain.ai.ChatMessages;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

@Component
public class MongoChatMemoryStore implements ChatMemoryStore {

   @Autowired
   private MongoTemplate mongoTemplate;

   public List<ChatMessages> getMessagesByUserId(Integer userId) {
       Criteria criteria = Criteria.where("userId").is(userId);
       Query query = new Query(criteria);
       List<ChatMessages> chatMessages = mongoTemplate.find(query, ChatMessages.class);
       if(chatMessages.isEmpty())
           return new ArrayList<>();
       return chatMessages;
   }

   public String createNewMemory() {
       // 1. 获取当前用户ID
       Integer userId = UserHolder.getUser().getUserId();

       // 2. 生成唯一memoryId
       String memoryId = new ObjectId().toString();

       // 3. 创建完整文档对象
       ChatMessages newDoc = new ChatMessages();
       newDoc.setMemoryId(memoryId.toString());
       newDoc.setUserId(userId);

       // 5. 执行插入操作
       mongoTemplate.insert(newDoc);
       return memoryId;
   }

    public void appendContentToMemory(String memoryId, String message, String newContent) {
        // 1. 根据memoryId查找文档
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);

        if (chatMessages == null) {
            throw new RuntimeException("Chat memory not found for memoryId: " + memoryId);
        }

        // 2. 反序列化现有的content
        List<ChatMessage> existingMessages = ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());

        // 3. 创建新消息并添加到列表
        // 用户提问信息
        ChatMessage userMessage = new UserMessage(message);
        existingMessages.add(userMessage);
        // AI返回的内容
        ChatMessage aiMessage = new AiMessage(newContent);
        existingMessages.add(aiMessage);

        // 4. 序列化更新后的消息列表
        String updatedContent = ChatMessageSerializer.messagesToJson(existingMessages);

        // 5. 更新文档
        Update update = new Update();
        update.set("content", updatedContent);
        mongoTemplate.updateFirst(query, update, ChatMessages.class);
    }
   @Override
   public List<ChatMessage> getMessages(Object memoryId) {
       Criteria criteria = Criteria.where("memoryId").is(memoryId);
       Query query = new Query(criteria);
       ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
       if(chatMessages == null)
           return new LinkedList<>();
       return ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());
   }

   @Override
   public void updateMessages(Object memoryId, List<ChatMessage> messages) {

       // 创建新文档
      if (memoryId.equals("0")) {
           // 1. 获取当前用户ID
           Integer userId = UserHolder.getUser().getUserId();

           // 2. 生成唯一memoryId
           memoryId = new ObjectId().toString();

           // 3. 创建完整文档对象
           ChatMessages newDoc = new ChatMessages();
           newDoc.setMemoryId(memoryId.toString());
           newDoc.setUserId(userId);

           // 5. 执行插入操作
           mongoTemplate.insert(newDoc);
       }
      else {
          {
              Criteria criteria = Criteria.where("memoryId").is(memoryId);
              Query query = new Query(criteria);

              Update update = new Update();
              update.set("content", ChatMessageSerializer.messagesToJson(messages));
              // 仅当插入新文档时才设置 userId 字段（文档已存在时不修改 userId）
              Integer userId = UserHolder.getUser().getUserId();
              if (userId != null) {
                  update.setOnInsert("userId", userId);
              }
              //根据query条件能查询出文档，则修改文档；否则新增文档
              mongoTemplate.upsert(query, update, ChatMessages.class);
          }
      }

   }

   @Override
   public void deleteMessages(Object memoryId) {
       Criteria criteria = Criteria.where("memoryId").is(memoryId);
       Query query = new Query(criteria) ;
       mongoTemplate.remove(query, ChatMessages.class);
   }
}

